package problems.algorithm.prefix_tree;

/**
 * 参考：https://zhuanlan.zhihu.com/p/28891541
 * @author guanjinping
 * @since 2022-09-18
 */
public class Templat {
    // 以小写字母组成的单词组为例
    public static class TrieNode {
        // 单词的个数
        public int count;
        // 前缀的个数
        public int prefix;
        TrieNode[] nextNode = new TrieNode[26];

        public TrieNode() {
            count = 0;
            prefix = 0;
        }
    }

    // 插入一个单词
    public void insert(TrieNode root, String word) {
        if (root == null || word == null || word.length() < 1) {
            return;
        }

        char[] chs = word.toCharArray();
        int n = chs.length;
        int j;

        for (int i = 0; i < n; ++i) {
            j = chs[i] - 'a';
            if (root.nextNode[j] == null) {
                root.nextNode[j] = new TrieNode();
            }
            root = root.nextNode[j];
            ++root.prefix;
        }
        //以该节点结尾的单词数+1
        ++root.count;
    }

    // 查找单词是否存在，如果存在，则返回数量，否则返回-1
    public int search(TrieNode root, String word) {
        if (root == null || word == null || word.length() < 1) return -1;

        char[] chs = word.toCharArray();
        int k;

        for (char c : chs) {
            k = c - 'a';
            if (root.nextNode[k] == null) return -1;
            root = root.nextNode[k];
        }

        return root.count == 0 ? -1 : root.count;
    }

    // 查询以str为前缀的单词数量
    public int searchPrefix(TrieNode root, String str) {
        if (root == null || str == null || str.length() < 1) return 0;
        char[] chs = str.toCharArray();
        for (char c : chs) {
            if (root.nextNode[c - 'a'] == null) return 0;
            root = root.nextNode[c - 'a'];
        }
        return root.prefix;
    }
}
